#include "../include/SSS_matvec.h"
#include "../include/SSS_utils.cuh"
#include "../../Solve/include/SSS_SOLVE.h"
#include "../../Setup/include/SSS_SETUP.h"
#include <cuda_runtime_api.h>
#include "../include/SSS_AMG.h"

//#include <math.h>
extern double setup_time ; 
extern double solve_time ; 
extern double total_time ; 

double SpGEMM_count = 0;
double SpMV_count = 0 ;

SSS_RTN SSS_solver_amg(SSS_MAT *A, SSS_VEC *x, SSS_VEC *b, SSS_AMG_PARS *pars)
{

    

    double total_setup = 0;
    double total_solve = 0;
    double total_amg = 0;
    int nnz, m, n;
    SSS_RTN rtn;
    SSS_AMG_PARS npars;

    SSS_AMG mg;
    double AMG_start, AMG_end;
    double sumb;

    
    cudaError_t EE;
    //
    cudaStream_t stream = NULL;

    printf("begin init cublas  ! \n");

    //cublas
    cublasHandle_t cublas_handle_AMG;
    cublasStatus_t cublas_status_AMG;
    cublasCreate(&cublas_handle_AMG);
    
    printf("end init cublas  ! \n");


    //cusparse
    // cusparseHandle_t     handle = NULL;
    // cusparseMatDescr_t  descrA = 0;
    // cusparseCreate(&handle) ;

    //cusolver
    // cusolverDnHandle_t cusolverH = NULL;
    // cusolverDnCreate(&cusolverH);
    // cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking);
    // cusolverDnSetStream(cusolverH, stream);

    //在device上给b开空间
    printf("cudamalloc b  ! \n");

    EE = cudaMalloc((void **)&b->d_d,(b->n) * sizeof(double));
    if(EE != cudaSuccess)
    {
        printf("cudaMalloc b Failed \n");
    }

    printf("cudamemcpy b  ! \n");
    EE = cudaMemcpy(b->d_d, b->d, b->n * sizeof(double), cudaMemcpyHostToDevice);
    if(EE != cudaSuccess)
    {
        printf("cudaMemcpy b Failed \n");
    }
    
    // 同步函数
    //cudaDeviceSynchronize();
    
    /* 计算向量b的L2范数 */
    //sumb = SSS_blas_vec_norm2(b);

    printf("begin norm2  ! \n");
    cublas_status_AMG = cublasDnrm2(cublas_handle_AMG,b->n,b->d_d,1,&sumb);
    printf("end norm2  ! \n");

    

    if(cublas_status_AMG != CUBLAS_STATUS_SUCCESS)
    {
        printf("norm computation unsuccessfull\n");
    }
    
    //cudaDeviceSynchronize();

    
    nnz = A->num_nnzs;
    m = A->num_rows;
    n = A->num_cols;

    AMG_start = SSS_get_time();

    // check matrix data
    if (m != n) {
        printf("### ERROR: A is not a square matrix!\n");
    }

    if (nnz <= 0) {
        printf("### ERROR: A has no nonzero entries!\n");
    }
    

    printf("begin setup phase ! \n");
    // Step 1: AMG setup phase
    double setup1 = SSS_get_time();
    SSS_amg_setup(&mg, A, pars);
    double setup2 = SSS_get_time();

    total_setup+=setup2-setup1;


    // Step 2: AMG solve phase
    double solve1 = SSS_get_time();
    rtn = SSS_amg_solve_cuda(&mg, x, b);
    
    double solve2 = SSS_get_time();

    total_solve+=solve2-solve1;

    // clean-up memory
    //SSS_amg_data_destroy(&mg);

    // print out CPU time if needed
    AMG_end = SSS_get_time();
    total_amg = AMG_end - AMG_start;

    setup_time = total_setup;
    solve_time = total_solve;

    printf("---------- AMG Time : ----------\n");
    printf("AMG totally time: %g ms  \n", total_amg);
    printf("AMG Setup time  : %g ms ,Setup / Total = %.2lf %%\n", total_setup , total_setup/ total_amg);
    printf("AMG Solve time  : %g ms ,Solve / Total = %.2lf %%\n", total_solve , total_solve / total_amg);

    printf("\n");

    printf("---------- kernel count : ----------\n");
    printf("SpGEMM_count :  %g   \n", SpGEMM_count);
    printf("SpMV_count   :  %g   \n", SpMV_count);

    //printf("Jacobi_count: %g ms  \n", total_amg);
    printf("\n");
    

    
    return rtn;
}